import streamlit as st

from init_vanna import MyVanna
from customer_logging import get_logger
from data.user_data import VAN_DB_CONFIG

logger = get_logger("chat-")


@st.cache_resource
def get_vanna():
    logger.info("初始化图表生成大模型")
    my_vanna = MyVanna(api_key="9b1f2582ae0c44c036ba1fdabed75c7b.FsfpGGSXKU4EZdsf")
    return my_vanna.vanna


vn = get_vanna()


@st.cache_data(show_spinner="Generating sample questions ...", ttl=60 * 60)
def generate_questions_cached():
    return vn.generate_questions()


@st.cache_data(show_spinner="Generating SQL query ...", ttl=60 * 60)
def generate_sql_cached(question: str):
    return vn.generate_sql(question=question)


@st.cache_data(show_spinner="Running SQL query ...", ttl=60 * 60)
def run_sql_cached(sql: str):
    vn.connect_to_mysql(**VAN_DB_CONFIG)
    return vn.run_sql(sql=sql)


@st.cache_data(show_spinner="Generating Plotly code ...", ttl=60 * 60)
def generate_plotly_code_cached(question, sql, df):
    code = vn.generate_plotly_code(question=question, sql=sql, df=df)
    return code


@st.cache_data(show_spinner="Running Plotly code ...", ttl=60 * 60)
def generate_plot_cached(code, df):
    return vn.get_plotly_figure(plotly_code=code, df=df)


@st.cache_data(show_spinner="Generating followup questions ...", ttl=60 * 60)
def generate_followup_cached(question, df, sql):
    return vn.generate_followup_questions(question=question, df=df, sql=sql)


@st.cache_data(show_spinner="deleting", ttl=60 * 60)
def remove_training_data(id_type):
    vn.remove_training_data(id=id_type)
